import { Context } from './Context';
import type { Middleware, TypedKey } from './types';

export class ContextWithMiddleware extends Context {
  private middlewares: Middleware[] = [];

  use(middleware: Middleware): () => void {
    this.middlewares.push(middleware);
    return () => {
      this.middlewares = this.middlewares.filter((m) => m !== middleware);
    };
  }

  setData<T>(key: TypedKey<T>, value: T, immediate = true): void {
    if (this.middlewares.length === 0) {
      super.setData(key, value, immediate);
      return;
    }

    const executeMiddleware = (index: number, currentValue: unknown) => {
      if (index >= this.middlewares.length) {
        super.setData(key, currentValue as T, immediate);
        return;
      }

      const middleware = this.middlewares[index];
      try {
        middleware(key.key, currentValue, (newValue) => {
          executeMiddleware(index + 1, newValue);
        });
      } catch (error) {
        console.error(`Middleware error for key ${key.key}:`, error);
        executeMiddleware(index + 1, currentValue);
      }
    };

    executeMiddleware(0, value);
  }
}
